#pragma once

#include <memory>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <boost/asio.hpp>
#include <boost/log/trivial.hpp>
#include "config.hpp"

namespace net_channel {
namespace client {

class DataChannel : public std::enable_shared_from_this<DataChannel> {
public:
    DataChannel(boost::asio::io_context& io_context,
                const std::string& local_address,
                uint16_t local_port,
                uint16_t remote_port,
                const std::string& protocol)
        : io_context_(io_context)
        , local_socket_(io_context)
        , remote_socket_(io_context)
        , local_endpoint_(boost::asio::ip::make_address(local_address), local_port)
        , remote_port_(remote_port)
        , protocol_(protocol) {}

    void start() {
        connect_to_local_service();
    }

private:
    void connect_to_local_service() {
        local_socket_.async_connect(local_endpoint_,
            [this](const boost::system::error_code& ec) {
                if (!ec) {
                    BOOST_LOG_TRIVIAL(info) << "Connected to local service at "
                                            << local_endpoint_.address().to_string()
                                            << ":" << local_endpoint_.port();
                    start_forwarding();
                } else {
                    BOOST_LOG_TRIVIAL(error) << "Failed to connect to local service: "
                                            << ec.message();
                }
            });
    }

    void start_forwarding() {
        do_read_local();
        do_read_remote();
    }

    void do_read_local() {
        auto self = shared_from_this();
        local_socket_.async_read_some(boost::asio::buffer(local_buffer_),
            [this, self](boost::system::error_code ec, std::size_t length) {
                if (!ec) {
                    do_write_remote(length);
                } else {
                    BOOST_LOG_TRIVIAL(error) << "Local read failed: " << ec.message();
                }
            });
    }

    void do_write_remote(std::size_t length) {
        auto self = shared_from_this();
        boost::asio::async_write(remote_socket_,
            boost::asio::buffer(local_buffer_, length),
            [this, self](boost::system::error_code ec, std::size_t /*length*/) {
                if (!ec) {
                    do_read_local();
                } else {
                    BOOST_LOG_TRIVIAL(error) << "Remote write failed: " << ec.message();
                }
            });
    }

    void do_read_remote() {
        auto self = shared_from_this();
        remote_socket_.async_read_some(boost::asio::buffer(remote_buffer_),
            [this, self](boost::system::error_code ec, std::size_t length) {
                if (!ec) {
                    do_write_local(length);
                } else {
                    BOOST_LOG_TRIVIAL(error) << "Remote read failed: " << ec.message();
                }
            });
    }

    void do_write_local(std::size_t length) {
        auto self = shared_from_this();
        boost::asio::async_write(local_socket_,
            boost::asio::buffer(remote_buffer_, length),
            [this, self](boost::system::error_code ec, std::size_t /*length*/) {
                if (!ec) {
                    do_read_remote();
                } else {
                    BOOST_LOG_TRIVIAL(error) << "Local write failed: " << ec.message();
                }
            });
    }

    boost::asio::io_context& io_context_;
    boost::asio::ip::tcp::socket local_socket_;
    boost::asio::ip::tcp::socket remote_socket_;
    boost::asio::ip::tcp::endpoint local_endpoint_;
    uint16_t remote_port_;
    std::string protocol_;
    std::array<char, 8192> local_buffer_;
    std::array<char, 8192> remote_buffer_;
};

class ConnectionManager {
public:
    explicit ConnectionManager(boost::asio::io_context& io_context)
        : io_context_(io_context) {}

    void start_service(const Config::ServiceInfo& service_info) {
        auto channel = std::make_shared<DataChannel>(
            io_context_,
            service_info.local_address,
            service_info.local_port,
            service_info.remote_port,
            service_info.protocol);
        channels_.insert(channel);
        channel->start();
    }

private:
    boost::asio::io_context& io_context_;
    std::unordered_set<std::shared_ptr<DataChannel>> channels_;
};

} // namespace client
} // namespace net_channel